use anyhow::Error;
use std::cmp::Ordering;
use crate::{db::Db, frame::Frame, tools::optimized_norm};

#[derive(Debug)]
struct FilterCondition {
    key: String,
    val: String,
}

#[derive(Debug)]
pub struct Search {
    collection_name: String,
    query: Vec<f32>,
    filters: Vec<FilterCondition>,
    k: usize,
}

impl Search {
    pub fn parse_from_frame(frame: Frame) -> Result<Self, Error> {
        let mut args = frame.get_args().into_iter();
        
        // 跳过命令名称
        args.next().ok_or_else(|| Error::msg("Missing command"))?;

        // 获取集合名称
        let collection_name = match args.next() {
            Some(name) => name,
            None => {
                return Err(Error::msg("Missing collection name"));
            }
        };

        // 获取向量内容
        let mut query = Vec::new();
        while let Some(arg) = args.next() {
            if arg.eq_ignore_ascii_case("K") || arg.eq_ignore_ascii_case("FILTER") {
                args = vec![arg].into_iter().chain(args).collect::<Vec<_>>().into_iter();
                break;
            }
            
            let parse_result = arg.parse::<f32>();
            let num = match parse_result {
                Ok(n) => n,
                Err(_) => {
                    let error_message = format!("Invalid float value: {}", arg);
                    return Err(Error::msg(error_message));
                }
            };
            query.push(num);
        }

        if query.is_empty() {
            return Err(Error::msg("Query vector cannot be empty"));
        }

        // 获取查询范围
        let mut k = 10;
        if let Some(arg) = args.next() {
            if arg.eq_ignore_ascii_case("K") {
                k = args.next().ok_or_else(|| Error::msg("Missing K value"))?.parse::<usize>().map_err(|_| Error::msg("Invalid K value"))?;
            } else if arg.eq_ignore_ascii_case("FILTER") {
                args = vec![arg].into_iter().chain(args).collect::<Vec<_>>().into_iter();
            }
        }

        // 获取过滤参数
        let mut filters = Vec::new();
        if let Some(arg) = args.next() {
            if !arg.eq_ignore_ascii_case("FILTER") {
                return Err(Error::msg(format!("Unexpected argument: {}", arg)));
            }

            for arg in args {
                let split_result = arg.split_once('=');
                let (key, val) = match split_result {
                    Some((k, v)) => (k, v),
                    None => {
                        let error_message = format!("Invalid filter format: {}", arg);
                        return Err(Error::msg(error_message));
                    }
                };
                
                filters.push(FilterCondition {
                    key: key.to_string(),
                    val: val.to_string(),
                });
            }
        }

        Ok(Self {
            collection_name,
            query,
            filters,
            k,
        })
    }

    pub fn apply(self, db: &Db) -> Result<Frame, Error> {

        let collection = match db.collections.get(&self.collection_name) {
            Some(c) => c,
            None => {
                return Err(Error::msg(format!("Collection '{}' not found", self.collection_name)));
            }
        };

        if self.query.len() != collection.dimension {
            return Err(Error::msg(format!("Query dimension mismatch (expected {}, found {})", collection.dimension, self.query.len())));
        }
        
        let query_norm = optimized_norm(self.query.iter().copied());
        let mut result = Vec::new();

        for (id, vec_data) in collection.vectors.iter() {
            let mut is_match = true;
            for cond in self.filters.iter() {
                if vec_data.metadata.get(&cond.key) != Some(&cond.val) {
                    is_match = false;
                    break;
                }
            }
            if is_match {
                let dot_product = self.query.iter().zip(&vec_data.vector).map(|(a, b)| a * b).sum::<f32>();
                let similarity = if query_norm == 0.0 || vec_data.norm == 0.0 {
                    0.0
                } else {
                    dot_product / (query_norm * vec_data.norm)
                };

                result.push((id.clone(), similarity));
            }
        }

        // 余弦相似度排序
        result.sort_by(|a, b| 
            b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)
        );

        let output = result.into_iter()
            .take(self.k)
            .map(|(id, score)| Frame::Array(vec![
                Frame::BulkString(score.to_string().into()),
                Frame::BulkString(id.into())
            ]))
            .collect();

        Ok(Frame::Array(output))
    }
}